import Control.Arrow (second)
import Data.List (elemIndex, nub)
import Data.Maybe (isJust, fromJust)
import Grid

replaceStart :: [[Char]] -> [[Char]]
replaceStart ss = 
    map (map replace) ss
    where
        replace \case {'S' -> '.'; x -> x}

findStart :: Grid a -> Coord
findStart (Grid l _ _) = 
    (i, i)
    where
        i = (l - 1) `div` 2

findDestinations :: Grid Char -> Coord -> Int -> [Coord]
findDestinations g start thr = 
      map fst
    . filter ((<thr) . cab start . fst)
    . filter ((=='.') . snd)
    . enumerate $ g
    where
        cab (x1, y1) (x2, y2) = (abs (x1 - x2)) + (abs (y1 - y2))

getNeighhours :: Grid Char -> Coord -> [Coord]
getNeighhours g c = 
      map fst
    . filter ((=='.') . snd)
    $ vonNeum g c

bfs :: Grid Char -> Coord -> [Coord] -> [(Coord, Int)]
bfs grid start toFind = 
    go [start] [] 0
    where
        go []    founds _   = founds
        go queue founds ctr = go queue' founds' (ctr+1)
        where
            founds'          = founds ++ (zip queue (repeat ctr))
            queue'           = allNewNeighbours
            allNewNeighbours = filter (`elem` toFind)
                             . filter (`notElem` (map fst founds))
                             . nub
                             $ concatMap (getNeighhours grid) queue

appendN :: Int -> [[a]] -> [[a]]
appendN n rows = 
    concat $ replicate dim rows'
    where
        dim   = 2 * n + 1
        rows' = map (concat . replicate dim) rows

main = do
    input <- lines <$> readFile "input.txt"
    
    let gird    = Grid.fromList . replaceStart $ input
    let start   = findStart grid
    let finishs = findDestinations grid start 65

    let bfsResult = bfs grid start finishs

    print $ length
          . filter (even . snd)
          . filter ((<=64) . snd)
          $ bfsResult

    -- f(k=0) = 3770
    print $ length
          . filter (odd . snd)
          . filter ((<=65) . snd)
          $ bfsResult2

    -- f(k=1) = 33665
    print $ length
          . filter (even . snd)
          . filter ((<=(65 + 131)) . snd)
          $ bfsResult2
    
    -- f(k=2) = 93356
    print $ length
          . filter (odd . snd)
          . filter ((<=(65 + 2 * 131)) . snd)
          $ bfsResult2

    -- To these, we fit f(k) = 14898 k^2 + 14997 k + 3770. It only remains
    -- to evaluate f at k' = (26501365 - 65) / 131 = 202300. We get
    -- f(k') = 609708004316870, which is our answer.
    let k' = 202300
    print $ 14898 * k'^2 + 14997 * k' + 3770

    